from langchain_chroma import Chroma
from langchain_core.documents import Document
from langchain_openai import  OpenAIEmbeddings
from langchain_core.runnables import RunnableLambda
import os

root_dir = os.path.dirname(os.path.abspath(__file__))

# 准备测试数据，假设我们提供的文档数据如下:
documents = [
    Document(
        page_content="狗是伟大的伴侣，以其忠诚和友好而闻名。", # 文档内容
        metadata={"source":"哺乳动物宠物文档"} # 文档元数据
    ),
    Document(
        page_content="猫是独立的宠物，通常喜欢自己的空间。",
        metadata={"source":"哺乳动物宠物文档"}
    ),
    Document(
        page_content="金鱼实时初学者的流行宠物，需要相对简单的护理。",
        metadata={"source": "鱼类宠物文档"}
    ),
    Document(
        page_content="白头鹰是犀利的猎手，视力很强大。",
        metadata={"source": "鸟类宠物文档"}
    )
]

# 创建一个向量数据库
vector_store = Chroma(
    collection_name="chroma_collection",
    embedding_function=OpenAIEmbeddings(model="text-embedding-3-large"),
    persist_directory=root_dir +"/vector_store/chroma_langchain_db",  # 指定持久化的目录，如果不存在该目录，则创建该目录；如果不设置该参数，就默认使用内存存储。
)

# 将文档添加到向量数据库中
vector_store.add_documents(documents)
# 相似度查询,返回相似的文档
retriever = RunnableLambda(vector_store.similarity_search).bind(k=4)
print(retriever.batch(['加菲猫', '鲸鱼']))

# 获取项目的根目录
